import cv2
import numpy as np
from PIL import Image
import matplotlib.pyplot as plt


def adaptive_histogram_equalization1(image_path, block_size=16):
    # 读取彩色图像
    img = Image.open(image_path).convert("RGB")
    img.show()
    if img is None:
        raise ValueError("图像路径错误或图像无法读取")

    # 转换到HSV颜色空间
    hsv_img = img.convert("HSV")
    h, s, v = hsv_img.split()

    # 将V通道（亮度）转换为NumPy数组进行处理
    v_array = np.array(v)
    height, width = v_array.shape

    # 创建与原图大小相同的结果数组
    v_eq_array = np.zeros_like(v_array, dtype=np.uint8)

    # 计算块的数量（考虑边界处理）
    blocks_y = height // block_size + (1 if height % block_size != 0 else 0)
    blocks_x = width // block_size + (1 if width % block_size != 0 else 0)

    # 为每个块计算局部直方图并应用均衡化
    for y in range(blocks_y):
        for x in range(blocks_x):
            # 计算当前块的边界
            y_start = y * block_size
            y_end = min(y_start + block_size, height)
            x_start = x * block_size
            x_end = min(x_start + block_size, width)

            # 提取当前块
            block = v_array[y_start:y_end, x_start:x_end]

            # 计算块的直方图
            hist, _ = np.histogram(block.flatten(), 256, [0, 256])

            # 计算累积分布函数（CDF）
            cdf = hist.cumsum()

            # 直方图均衡化映射
            cdf_m = np.ma.masked_equal(cdf, 0)  # 避免除以零
            cdf_m = (cdf_m - cdf_m.min()) * 255 / (cdf_m.max() - cdf_m.min())
            cdf = np.ma.filled(cdf_m, 0).astype('uint8')

            # 应用映射到当前块
            v_eq_array[y_start:y_end, x_start:x_end] = cdf[block]

    # 将处理后的V通道转换回PIL图像
    v_eq = Image.fromarray(v_eq_array)

    # 合并处理后的V通道与原始H、S通道
    hsv_eq = Image.merge("HSV", (h, s, v_eq))

    # 转回RGB空间
    rgb_eq = hsv_eq.convert("RGB")
    rgb_eq.show()

    return rgb_eq



def adaptive_histogram_equalization(image_path, block_size=16, clip_limit=None):
    """
    使用OpenCV实现自适应直方图均衡化（AHE）

    参数:
    - image_path: 输入图像路径
    - block_size: 分块大小（默认16x16）
    - clip_limit: 对比度限制（默认为None，即普通AHE；若指定值则变为CLAHE）
    """
    # 读取图像并转换为灰度图
    img = cv2.imread(image_path)
    if img is None:
        raise ValueError(f"无法读取图像: {image_path}")

    # 转换为灰度图
    gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY)

    # 获取图像尺寸
    height, width = gray.shape

    # 创建与原图相同大小的结果数组
    result = np.zeros_like(gray, dtype=np.uint8)

    # 计算分块数量
    blocks_x = (width + block_size - 1) // block_size
    blocks_y = (height + block_size - 1) // block_size

    # 为每个分块创建映射表
    maps = np.zeros((blocks_y, blocks_x, 256), dtype=np.uint8)

    # 遍历每个分块
    for y in range(blocks_y):
        for x in range(blocks_x):
            # 计算当前分块的边界
            x_start = min(x * block_size, width - 1)
            x_end = min((x + 1) * block_size, width)
            y_start = min(y * block_size, height - 1)
            y_end = min((y + 1) * block_size, height)

            # 提取分块区域
            block = gray[y_start:y_end, x_start:x_end]

            # 计算直方图
            hist = cv2.calcHist([block], [0], None, [256], [0, 256])

            # 如果指定了clip_limit，则进行对比度限制（即CLAHE）
            if clip_limit is not None:
                # 裁剪直方图
                hist = np.clip(hist, 0, clip_limit * block.size / 256)
                # 重新分配超出的部分
                excess = np.sum(hist - clip_limit * block.size / 256)
                hist += excess / 256
                hist = np.clip(hist, 0, None)  # 确保不会出现负值

            # 计算累积分布函数(CDF)
            cdf = hist.cumsum()

            # 归一化CDF到0-255范围
            cdf_normalized = ((cdf - cdf[0]) * 255 / (cdf[-1] - cdf[0])).astype(np.uint8)

            # 保存映射表
            maps[y, x] = cdf_normalized.flatten()

    # 对每个像素进行双线性插值
    """for y in range(height):
        for x in range(width):
            # 计算当前像素所在的分块坐标(浮点值)
            fx = x / block_size
            fy = y / block_size

            # 确定相邻的四个分块
            x1 = int(fx)
            y1 = int(fy)
            x2 = min(x1 + 1, blocks_x - 1)
            y2 = min(y1 + 1, blocks_y - 1)

            # 计算权重
            wx2 = fx - x1
            wy2 = fy - y1
            wx1 = 1 - wx2
            wy1 = 1 - wy2

            # 获取当前像素值
            pixel_value = gray[y, x]

            # 从四个相邻分块获取映射值
            v11 = maps[y1, x1, pixel_value]
            v21 = maps[y1, x2, pixel_value]
            v12 = maps[y2, x1, pixel_value]
            v22 = maps[y2, x2, pixel_value]

            # 双线性插值
            result[y, x] = np.clip(
                v11 * wx1 * wy1 + v21 * wx2 * wy1 + v12 * wx1 * wy2 + v22 * wx2 * wy2,
                0, 255
            ).astype(np.uint8)"""

    # 将处理后的灰度图转回彩色显示
    result_color = cv2.cvtColor(maps, cv2.COLOR_GRAY2BGR)
    # 显示原始图像和处理后的图像
    plt.figure(figsize=(12, 6))
    plt.subplot(1, 2, 1)
    plt.imshow(cv2.cvtColor(img, cv2.COLOR_BGR2RGB))
    plt.title('原始图像')
    plt.axis('off')
    plt.subplot(1, 2, 2)
    plt.imshow(cv2.cvtColor(result_color, cv2.COLOR_BGR2RGB))
    plt.title(f'AHE增强 (块大小={block_size}x{block_size})')
    plt.axis('off')
    plt.tight_layout()
    plt.show()
    return result_color
if __name__ == "__main__":
    image_path = "D:\jvnhenghua\\test.jpg"
    adaptive_histogram_equalization(image_path)
#这个原版AHE不经过一些特殊处理的话效果差的像依托屎
#这个代码不是我本人写的我也看不懂，AHE其实也没学习的必要，有它的上位替代CLAHE